# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import logging
import os
from concurrent.futures import ThreadPoolExecutor

import aiofiles
from tornado.concurrent import run_on_executor
from tornado.ioloop import IOLoop

from coin.base.timestamp import get_timestamp
from coin.flow.topic_record import TopicRecordReader


class _AsyncTailReaderBase:
  def __init__(self, path, read_from_tail=True, log_callback=None, ioloop=None):
    self._logger = logging.getLogger(__name__)

    self._path = path
    self._read_from_tail = read_from_tail
    self._ioloop = ioloop or IOLoop.current()

    self._last_modified_ts = None
    self._update_last_modified_ts()
    self._log_callback = log_callback

  def __del__(self):
    self.close()

  def close(self):
    pass

  @property
  def last_modified_ts(self):
    return self._last_modified_ts

  def _update_last_modified_ts(self):
    self._last_modified_ts = get_timestamp()

  async def handle_modified_async(self):
    raise NotImplementedError()

  def handle_modified(self):
    self._ioloop.add_callback(self.handle_modified_async)


class AsyncLogTailReader(_AsyncTailReaderBase):
  def __init__(self, *args, chunk_size=512 * 1024, **kwargs):
    super().__init__(*args, **kwargs)

    self._last_pos = os.path.getsize(self._path) if self._read_from_tail else 0
    self._file = None
    self._chunk_size = chunk_size
    self._buffer = ''
    self._running = False
    self._modified = False

  def close(self):
    if self._file:
      self._file.close()
      self._file = None

  async def handle_modified_async(self):
    if self._running:
      self._modified = True
      return
    try:
      self._running = True
      self._modified = False
      if not self._file:
        self._file = await aiofiles.open(self._path, 'r')
        await self._file.seek(self._last_pos)

      file_size = os.path.getsize(self._path)
      if self._last_pos > file_size:
        # Truncated.
        self._last_pos = file_size
        await self._file.seek(self._last_pos)
      self._buffer += await self._file.read()
      self._last_pos = await self._file.tell()
      if not self._buffer:
        return
      length = self._buffer.rfind('\n', 0, self._chunk_size) + 1
      if len(self._buffer) >= self._chunk_size:
        self._modified = True
        if length == 0:
          length = self._chunk_size
      if length > 0:
        data = self._buffer[:length]
        self._buffer = self._buffer[length:]
        self._log_callback(self._path, data)
        self._update_last_modified_ts()
    finally:
      self._running = False
      if self._modified:
        self._ioloop.add_callback(self.handle_modified_async)


class AsyncProtoLogTailReader(_AsyncTailReaderBase):
  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self._reader = None

    self._running = False
    self._modified = False
    self._executor = None

  @run_on_executor(executor='_executor')
  def _read_last_data(self):
    last_data = None
    while True:
      data = self._reader.read()
      if data is None:
        break
      else:
        last_data = data
    return last_data

  @run_on_executor(executor='_executor')
  def _read_chunk(self, max_chunk_size=100):
    res = []
    eof_reached = False
    while True:
      data = self._reader.read()
      if data is None:
        eof_reached = True
        break
      else:
        res.append(data)
        if len(res) >= max_chunk_size:
          break
    return (eof_reached, res)

  def __del__(self):
    self._logger.info('%s cleared.', self._path)
    self.close()

  def close(self):
    self._logger.info('%s closed.', self._path)
    if self._reader is not None:
      self._reader.close()
      self._reader = None
    if self._executor is not None:
      self._executor.shutdown(wait=False)
      self._executor = None

  async def _prepare_reader(self):
    self._executor = ThreadPoolExecutor(max_workers=1)
    self._reader = TopicRecordReader(self._path)
    if self._read_from_tail:
      last_data = await self._read_last_data()
      self._logger.info('%s(%s) added.', self._path, self._read_from_tail)
      if last_data is not None:
        return [last_data]
    return []

  async def handle_modified_async(self):
    if self._running:
      self._modified = True
      return
    try:
      self._running = True
      self._modified = False

      if self._reader is None:
        eof_reached = True
        res = await self._prepare_reader()
      else:
        eof_reached, res = await self._read_chunk()
      for data in res:
        self._log_callback(self._path, data)
      if res:
        self._update_last_modified_ts()
      if not eof_reached:
        self._modified = True
    finally:
      self._running = False
      if self._modified:
        self._ioloop.add_callback(self.handle_modified_async)
